import base64
import os
import time

import cv2
import flask
import numpy as np
from PIL import Image
from flask import Flask, request, redirect, Response, jsonify

from yolo import YOLO

app = Flask(__name__)


@app.route('/')
def index():
    return flask.render_template("index.html")


@app.route('/photo')
def photo():
    return flask.render_template("photo.html")


@app.route('/now')
def now():
    return flask.render_template("now.html")


@app.route("/video")
def video():
    return flask.render_template("video.html")


@app.route("/video_pre")
def video_pre():
    return flask.render_template("video_pre.html")


@app.route("/plane")
def plane():
    return flask.render_template("plane.html")


@app.route("/plane_show")
def plane_show():
    return flask.render_template("plane_show.html")


yolo = YOLO()
# ----------------------------------------------------------------------------------------------------------#
#   video_path          用于指定视频的路径，当video_path=0时表示检测摄像头
#                       想要检测视频，则设置如video_path = "xxx.mp4"即可，代表读取出根目录下的xxx.mp4文件。
#   video_save_path     表示视频保存的路径，当video_save_path=""时表示不保存
#                       想要保存视频，则设置如video_save_path = "yyy.mp4"即可，代表保存为根目录下的yyy.mp4文件。
#   video_fps           用于保存的视频的fps
#
#   video_path、video_save_path和video_fps仅在mode='video'时有效
#   保存视频时需要ctrl+c退出或者运行到最后一帧才会完成完整的保存步骤。4
# ----------------------------------------------------------------------------------------------------------#
video_path = "rtmp://192.168.123.97/live"   #"rtmp://192.168.123.96/live"
# video_path = 0
video_save_path = ''
video_fps = 25.0


# 实时检测
class Camera(object):
    def __init__(self):
        self.capture = cv2.VideoCapture(video_path)

    def __del__(self):
        self.capture.release()

    def live(self):
        ref, image = self.capture.read()
        if not ref:
            raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")
        fps = 0.0
        while True:
            t1 = time.time()
            # 读取某一帧
            ref, frame = self.capture.read()
            if not ref:
                break
            # 格式转变，BGRtoRGB
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            # 转变成Image
            frame = Image.fromarray(np.uint8(frame))
            # 进行检测
            frame = np.array(yolo.detect_image(frame))
            # RGBtoBGR满足opencv显示格式
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

            fps = (fps + (1. / (time.time() - t1))) / 2
            print("fps= %.2f" % fps)
            # 在图像上画图
            frame = cv2.putText(frame, "fps= %.2f" % fps, (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            # 图像实时展示
            cv2.imshow("video", frame)
            return frame

    def get_frame(self):
        image = self.live()
        ret, jpeg = cv2.imencode('.jpg', image)
        return jpeg.tobytes()


def gen(camera):
    while True:
        frame = camera.get_frame()
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')


@app.route('/now_feed')
def now_feed():
    return Response(gen(Camera()),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


# 图片检测
def photo_pre(img):
    while True:
        try:
            image = Image.open(img)
        except:
            print('Open Error! Try again!')
            continue
        else:
            r_image = yolo.detect_image(image)
            return r_image


@app.route("/photo_feed", methods=["POST"])
def photo_feed():
    if request.method == "POST":
        photos = request.files['photo']
        img_name = "pre.jpg"
        photo_path = os.path.join('img/', img_name)
        print(photo_path)
        if not photo_path:
            os.mkdir(photo_path)
        photos.save(photo_path)
        time.sleep(5)
        r_img = photo_pre('img/pre.jpg')
        r_img.save(photo_path)
        image = open("img/pre.jpg", "rb")
        data = base64.b64encode(image.read())
        data = str(data, encoding="utf-8")
        data = "data:image/jpg;base64," + data
        return data


# 视频检测


class Camera_video(object):
    def __init__(self, video_path):
        self.capture = cv2.VideoCapture(video_path)

    def __del__(self):
        self.capture.release()

    def live(self):
        ref, image = self.capture.read()
        if not ref:
            raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")
        fps = 0.0
        while True:
            t1 = time.time()
            # 读取某一帧
            ref, frame = self.capture.read()
            if not ref:
                break
            # 格式转变，BGRtoRGB
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            # 转变成Image
            frame = Image.fromarray(np.uint8(frame))
            # 进行检测
            frame = np.array(yolo.detect_image(frame))
            # RGBtoBGR满足opencv显示格式
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

            fps = (fps + (1. / (time.time() - t1))) / 2
            print("fps= %.2f" % fps)
            # 在图像上画图
            frame = cv2.putText(frame, "fps= %.2f" % fps, (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            # 图像实时展示
            cv2.imshow("video", frame)
            return frame

    def get_frame_video(self):
        image = self.live()
        ret, jpeg = cv2.imencode('.jpg', image)
        return jpeg.tobytes()


def gen_video(camera):
    while True:
        frame = camera.get_frame_video()
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')


@app.route('/video_feed', methods=["POST"])
def video_feed():
    video_file = request.files['video']
    video_name = video_file.filename
    video_path = os.path.join("video/" + video_name)
    video_file.save(video_path)

    return Response(gen_video(Camera_video(video_path)),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


# 无人机实时画面
class Camera_plane(object):
    def __init__(self):
        self.capture = cv2.VideoCapture(video_path)

    def __del__(self):
        self.capture.release()

    def live(self):
        ref, image = self.capture.read()
        if not ref:
            raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")
        fps = 0.0
        while True:
            t1 = time.time()
            # 读取某一帧
            ref, frame = self.capture.read()
            if not ref:
                break

            fps = (fps + (1. / (time.time() - t1))) / 2
            print("fps= %.2f" % fps)
            # 在图像上画图
            frame = cv2.putText(frame, "fps= %.2f" % fps, (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            # 图像实时展示
            cv2.imshow("video", frame)
            return frame

    def get_frame_plane(self):
        image = self.live()
        ret, jpeg = cv2.imencode('.jpg', image)
        return jpeg.tobytes()


def gen_plane(camera):
    while True:
        frame = camera.get_frame_plane()
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')


@app.route('/plane_feed')
def plane_feed():
    return Response(gen_plane(Camera_plane()),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


if __name__ == "__main__":
    app.run(port=5000, debug=True)
